{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:38:20.588780Z",
     "start_time": "2019-05-26T01:38:20.579169Z"
    },
    "code_folding": [],
    "run_control": {
     "marked": false
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "from Util import my_helper\n",
    "import copy\n",
    "import itertools\n",
    "import random\n",
    "import pickle as cPickle\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.layers import core as core_layers\n",
    "\n",
    "\n",
    "from Model import LM\n",
    "from Util.myAttWrapper import SelfAttWrapper\n",
    "from Util import myResidualCell\n",
    "from Util.bleu import BLEU\n",
    "from Util.myUtil import *\n",
    "\n",
    "\n",
    "tf.logging.set_verbosity(tf.logging.INFO)\n",
    "sess_conf = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:38:21.877344Z",
     "start_time": "2019-05-26T01:38:21.834864Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def _construct_blank(num, length, style='random'):\n",
    "    data = X_indices[:num]\n",
    "    blank_data = []\n",
    "    for d in data:\n",
    "        if style == 'random':\n",
    "            pos_list = np.sort(random.sample(range(len(d)-2), length)).tolist()\n",
    "            blank_data.append((d, pos_list))\n",
    "        elif style == 'middle':\n",
    "            l = int(length * (len(d) - 1))\n",
    "            pos_list  = [int((len(d)-1-l) / 2.0) + i for i in range(l)]\n",
    "            blank_data.append((d, pos_list))\n",
    "    return blank_data\n",
    "\n",
    "def show_blank(idx, pos):\n",
    "    t = [len(id2w)+5 if i in pos else a for i,a in enumerate(idx)]\n",
    "    s = ' '.join([id2w.get(tt,'_') for tt in t])\n",
    "    return s\n",
    "\n",
    "def idx2str(idx):\n",
    "    return \" \".join(id2w.get(idxx, '_') for idxx in idx)\n",
    "\n",
    "def str2idx(idx):\n",
    "    idx = idx.strip()\n",
    "    return [w2id[idxx] for idxx in idx.split(' ')]\n",
    "\n",
    "def cal_candidate_list(inputs, pos):\n",
    "    idx = copy.deepcopy(inputs)\n",
    "    idx[pos] = model.qid_list[0]\n",
    "    candidate_list = []\n",
    "    for t in range(len(model.dp.X_w2id)):\n",
    "        temp = copy.deepcopy(idx)\n",
    "        temp = [k if k!=model.qid_list[0] else t for k in temp]\n",
    "        candidate_list.append(temp)\n",
    "    return candidate_list\n",
    "\n",
    "\n",
    "def replace_list(idx, pos_list, target):\n",
    "    t = [idxx for idxx in idx]\n",
    "    if target:\n",
    "        for i,p in enumerate(pos_list):\n",
    "            t[p] = target[i]\n",
    "    else:\n",
    "        for i,p in enumerate(pos_list):\n",
    "            t[p] = -1\n",
    "    return t\n",
    "\n",
    "def cal_optimal(idx, pos, max_it=-1):\n",
    "    X_batch = cal_candidate_list(idx, pos)\n",
    "    X_batch_len = [len(x) for x in X_batch]\n",
    "    \n",
    "    if max_it > 0:\n",
    "        batch_loss = []\n",
    "        t = 0\n",
    "        while t+max_it<len(X_batch):\n",
    "            batch_loss += model.sess.run(model.batch_loss, {model.X: X_batch[t:t+max_it],\n",
    "                                                        model.X_seq_len: X_batch_len[t:t+max_it],\n",
    "                                                        model.output_keep_prob:1,\n",
    "                                                        model.input_keep_prob:1}).tolist()\n",
    "            t += max_it\n",
    "        batch_loss += model.sess.run(model.batch_loss, {model.X: X_batch[t:],\n",
    "                                                        model.X_seq_len: X_batch_len[t:],\n",
    "                                                        model.output_keep_prob:1,\n",
    "                                                        model.input_keep_prob:1}).tolist()\n",
    "    else:\n",
    "        batch_loss = model.sess.run(model.batch_loss, {model.X: X_batch,\n",
    "                                                        model.X_seq_len: X_batch_len,\n",
    "                                                        model.output_keep_prob:1,\n",
    "                                                        model.input_keep_prob:1}).tolist()\n",
    "    argsort_batch_loss = np.argsort(batch_loss)\n",
    "    sort_loss = [batch_loss[i] for i in argsort_batch_loss]\n",
    "    sort_idx = [X_batch[i][pos] for i in argsort_batch_loss]\n",
    "    return sort_loss, sort_idx\n",
    "\n",
    "def cal_dist(vector1, vector2):\n",
    "    return np.linalg.norm(vector1-vector2), np.dot(vector1,vector2)/(np.linalg.norm(vector1)*(np.linalg.norm(vector2)))\n",
    "\n",
    "\n",
    "def _init_blank(idx, pos):\n",
    "    c_idx = copy.deepcopy(idx)\n",
    "    o_idx = []\n",
    "    pos = np.sort(pos).tolist()\n",
    "    for i in range(len(idx)):\n",
    "        if i in pos:\n",
    "            if i == 0:\n",
    "                o_idx.append(w2id[model.generate(1)[1][0][0]])\n",
    "            else:\n",
    "                prefix = idx2str(o_idx)\n",
    "                infer = model.infer(prefix)\n",
    "                infer = infer[np.argmax([len(inf) for inf in infer])]\n",
    "                if len(str2idx(infer)) <= i:\n",
    "                    o_idx.append(w2id['<PAD>'])\n",
    "                else:\n",
    "                    o_idx.append(str2idx(infer)[i])\n",
    "        else:\n",
    "            o_idx.append(c_idx[i])\n",
    "    init_word = [id2w[o_idx[i]] for i in pos]    \n",
    "    return o_idx, idx2str(o_idx), init_word\n",
    "\n",
    "def _init_data(name):\n",
    "    w2id, id2w = cPickle.load(open('Data/%s/w2id_id2w.pkl' % name,'rb'))\n",
    "    X_indices = cPickle.load(open('Data/%s/index.pkl' % name,'rb'))\n",
    "    return X_indices, w2id, id2w\n",
    "\n",
    "def _init_model(name, lr=10.0, l1_reg_lambda=0.00, l2_reg_lambda=0.00, close_loss_rate=0.00):\n",
    "    qid_list = cPickle.load(open('Data/%s/qid_list.pkl'%name,'rb'))\n",
    "    qid_list = [w2id[w] for w in qid_list]\n",
    "    rnn_size = dict()\n",
    "    rnn_size['Poem'] =  512\n",
    "    rnn_size['Daily'] = 512 \n",
    "    rnn_size['APRC'] = 1024\n",
    "    \n",
    "    num_layer = dict()\n",
    "    num_layer['Poem'] = 2\n",
    "    num_layer['Daily'] = 1\n",
    "    num_layer['APRC'] = 1\n",
    "    \n",
    "    max_infer_length = dict()\n",
    "    max_infer_length['Poem'] = 33\n",
    "    max_infer_length['Daily'] = 50\n",
    "    max_infer_length['APRC'] = 36\n",
    "    \n",
    "    model_iter = dict()\n",
    "    model_iter['Poem'] = 30\n",
    "    model_iter['Daily'] = 30 \n",
    "    model_iter['APRC'] = 20\n",
    "    \n",
    "    assert name in ['Poem','Daily', 'APRC']\n",
    "\n",
    "    BATCH_SIZE = 256\n",
    "    NUM_EPOCH = 30\n",
    "    train_dir ='Model/%s' % name\n",
    "    dp = LM_DP(X_indices, w2id, BATCH_SIZE, n_epoch=NUM_EPOCH)\n",
    "    g = tf.Graph() \n",
    "    sess = tf.Session(graph=g, config=sess_conf) \n",
    "    with sess.as_default():\n",
    "        with sess.graph.as_default():\n",
    "            model = LM(\n",
    "                dp = dp,\n",
    "                rnn_size = rnn_size[name],\n",
    "                n_layers = num_layer[name],\n",
    "                decoder_embedding_dim = rnn_size[name],\n",
    "                cell_type='lstm',\n",
    "                close_loss_rate = close_loss_rate,\n",
    "                max_infer_length = max_infer_length[name],\n",
    "                att_type='B',\n",
    "                qid_list = qid_list,\n",
    "                lr = lr,\n",
    "                l1_reg_lambda = l1_reg_lambda,\n",
    "                l2_reg_lambda = l2_reg_lambda,\n",
    "                is_save = False,\n",
    "                residual = True,\n",
    "                is_jieba = False,\n",
    "                sess=sess\n",
    "            )\n",
    "\n",
    "\n",
    "    util = LM_util(dp=dp, model=model)\n",
    "    model.restore('Model/%s/model-%d'% (name,model_iter[name])) # restore pre-train model\n",
    "    return model\n",
    "\n",
    "\n",
    "\n",
    "def _reload(name):\n",
    "    rnn_size = dict()\n",
    "    rnn_size['Poem'] =  512\n",
    "    rnn_size['Daily'] = 512 \n",
    "    rnn_size['APRC'] = 1024\n",
    "    \n",
    "    num_layer = dict()\n",
    "    num_layer['Poem'] = 2\n",
    "    num_layer['Daily'] = 1\n",
    "    num_layer['APRC'] = 1\n",
    "    \n",
    "    max_infer_length = dict()\n",
    "    max_infer_length['Poem'] = 33\n",
    "    max_infer_length['Daily'] = 50\n",
    "    max_infer_length['APRC'] = 36\n",
    "    \n",
    "    model_iter = dict()\n",
    "    model_iter['Poem'] = 30\n",
    "    model_iter['Daily'] = 30 \n",
    "    model_iter['APRC'] = 20\n",
    "    \n",
    "    assert name in ['Poem','Daily', 'APRC']\n",
    "\n",
    "    model.restore('Model/%s/model-%d'% (name,model_iter[name]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TIGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:38:22.745667Z",
     "start_time": "2019-05-26T01:38:22.718086Z"
    }
   },
   "outputs": [],
   "source": [
    "def cal_optimizer_gibbs(inputs, pos_list, init_word, name='Nesterov', epoch=10, is_show=True, top_k=100, upper_size=2, distance='l2'):\n",
    "    total_tic = time.time()\n",
    "    init_time = 0.0\n",
    "    update_time = 0.0\n",
    "    assign_time = 0.0\n",
    "    search_time = 0.0\n",
    "    cal_pre_time = 0.0\n",
    "    cal_next_time = 0.0\n",
    "    \n",
    "    tic = time.time()\n",
    "    upper_cnt = 0\n",
    "    \n",
    "    # prepare\n",
    "    K = len(pos_list)\n",
    "    assert K <= len(model.qid_list)\n",
    "    idx = copy.deepcopy(inputs)\n",
    "    for i,pos in enumerate(pos_list):\n",
    "        idx[pos] = model.qid_list[i]\n",
    "    if init_word:\n",
    "        pre_p_list = [w2id[t] for t in init_word]\n",
    "    else:\n",
    "        pre_p_list = model.qid_list[:K]\n",
    "    next_p_list = []\n",
    "    pre_sentence = replace_list(idx, pos_list, pre_p_list)\n",
    "    epoch_sentence  = replace_list(idx, pos_list, pre_p_list)\n",
    "    word_emb = model.sess.run(model.decoder_embedding)\n",
    "        \n",
    "    # init specific embedding \n",
    "    if init_word:\n",
    "        feed_dict = dict()\n",
    "        for j in range(K):\n",
    "            feed_dict[model.assgin_placeholder_list[j]] = word_emb[[w2id[init_word[j]]]]\n",
    "        model.sess.run(model.assign_op_list[:K],feed_dict)\n",
    "    init_time += time.time()-tic\n",
    "\n",
    "    # search\n",
    "    for i in range(epoch):\n",
    "        if i > 1 and epoch_sentence == replace_list(idx, pos_list, pre_p_list):\n",
    "            upper_cnt += 1\n",
    "            if upper_cnt >= upper_size:\n",
    "                if is_show:\n",
    "                    print('total_epoch %d'% (i+1))\n",
    "                break\n",
    "        else:\n",
    "            upper_cnt = 0\n",
    "        epoch_sentence = replace_list(idx, pos_list, pre_p_list)\n",
    "        if is_show:\n",
    "            print('epoch %d :' %(i+1), idx2str(epoch_sentence))\n",
    "        ep_tic = time.time()    \n",
    "        for k in range(K):\n",
    "            # O-step\n",
    "            pre_sentence = replace_list(idx, pos_list, pre_p_list)\n",
    "            tic = time.time()\n",
    "            if distance == 'cos':\n",
    "                v, o = model.sess.run([model.nearby_val, model.nearby_idx], {model.nearby_word:[model.qid_list[k]]})\n",
    "                nearset = o[0][1]\n",
    "            else:\n",
    "                v, o = model.sess.run([model.eu_nearby_val, model.eu_nearby_idx], {model.nearby_word:[model.qid_list[k]]})\n",
    "                nearset = o[1]\n",
    "            loss, _ = model.sess.run([model.update_loss, \n",
    "                                      model.update_op[name+'_%d' % k]], \n",
    "                                         {model.X: [idx], \n",
    "                                          model.X_seq_len: [len(idx)], \n",
    "                                          model.Y:[pre_sentence],\n",
    "                                          model.output_keep_prob:1,\n",
    "                                          model.input_keep_prob:1,\n",
    "                                          model.nearest_emb_placeholder:word_emb[[nearset]]})\n",
    "            update_time += time.time() - tic\n",
    "\n",
    "\n",
    "            # P-step\n",
    "            if i % 1 == 0:\n",
    "                # candidate\n",
    "                tic = time.time()\n",
    "                if distance == 'cos':\n",
    "                    v, o = model.sess.run([model.nearby_val, model.nearby_idx], {model.nearby_word:[model.qid_list[k]]})\n",
    "                    candi_pos = o[0][1:top_k+1].tolist() + [pre_p_list[k]]\n",
    "                else:\n",
    "                    v, o = model.sess.run([model.eu_nearby_val, model.eu_nearby_idx], {model.nearby_word:[model.qid_list[k]]})\n",
    "                    candi_pos = o[1:top_k+1].tolist() + [pre_p_list[k]]\n",
    "                    \n",
    "                candi_list = [[pre_p_list[j] if j!=k else t for j in range(len(pre_p_list))] for t in candi_pos]\n",
    "                next_sentences = [replace_list(idx, pos_list, candi) for candi in candi_list]\n",
    "                search_time += time.time() - tic\n",
    "                \n",
    "                # cal loss\n",
    "                tic = time.time()\n",
    "                next_loss_list = model.sess.run(model.batch_loss, {model.X: next_sentences, \n",
    "                                        model.X_seq_len: [len(idx) for j in range(len(next_sentences))], \n",
    "                                        model.output_keep_prob:1,\n",
    "                                        model.input_keep_prob:1})\n",
    "                argmin_idx = np.argmin(next_loss_list)\n",
    "                next_p_pos = candi_pos[argmin_idx]\n",
    "                cal_next_time += time.time()-tic\n",
    "                # update\n",
    "                tic = time.time()\n",
    "                if next_p_pos != pre_p_list[k]:\n",
    "                    model.sess.run(model.assign_op_list[k],{model.assgin_placeholder_list[k]:word_emb[[next_p_pos]]})\n",
    "                    pre_p_list[k] = next_p_pos\n",
    "                    \n",
    "                assign_time += time.time() - tic\n",
    "                if is_show:\n",
    "                    print('epoch %d_%d :' % (i+1, k),idx2str(replace_list(idx, pos_list, pre_p_list)))\n",
    "                \n",
    "    tic = time.time()    \n",
    "    pre_sentence = replace_list(idx, pos_list, pre_p_list)\n",
    "    loss = model.sess.run(model.loss, {model.X: [pre_sentence], \n",
    "                        model.X_seq_len: [len(idx)], \n",
    "                        model.output_keep_prob:1,\n",
    "                        model.input_keep_prob:1})\n",
    "    cal_pre_time += time.time() - tic\n",
    "    total_time = time.time() - total_tic\n",
    "    \n",
    "    return pre_p_list, idx2str(pre_p_list),loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:38:23.739860Z",
     "start_time": "2019-05-26T01:38:23.434859Z"
    }
   },
   "outputs": [],
   "source": [
    "# initialize blank with left-to-right greedy beam search\n",
    "f_init = cPickle.load(open('results/_URNN-f_res.pkl','rb'))  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:42:54.580893Z",
     "start_time": "2019-05-26T01:42:01.554969Z"
    }
   },
   "outputs": [],
   "source": [
    "task_name = 'APRC' # 'Daily', 'Poem'\n",
    "assert task_name in ['Poem','Daily', 'APRC']\n",
    "X_indices, w2id, id2w = _init_data(task_name)\n",
    "model = _init_model(task_name, lr=10.0)\n",
    "\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-26T01:44:42.226194Z",
     "start_time": "2019-05-26T01:44:12.389464Z"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "is_init = True\n",
    "for length_ratio in [0.25, 0.5, 0.75]:\n",
    "    for style in ['random', 'middle']:\n",
    "        blank_data = cPickle.load(open('Data/%s/%d_%s.pkl'%(task_name, int(length_ratio*100), style),'rb'))\n",
    "        i = random.sample(range(5000), 1)[0]\n",
    "        idx, pos_list = blank_data[i]\n",
    "        prefix = '%s_%d_%s' % (task_name, int(length_ratio*100), style)\n",
    "        model._opt_init()\n",
    "        if is_init:\n",
    "            init_word = [id2w[f_init[prefix+'_URNN-f'][i][p]] for p in pos_list]\n",
    "\n",
    "        else:\n",
    "            init_word = random.sample(w2id.keys(), len(pos_list))\n",
    "        sid, sw, loss = cal_optimizer_gibbs(idx, pos_list, init_word = init_word, is_show=False)\n",
    "        print('Template:', show_blank(idx, pos_list))\n",
    "        print('GroundTruth:', idx2str(idx))\n",
    "        print('TIGS:', idx2str(replace_list(idx, pos_list, sid)))\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "165px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
